package db

import (
	"context"
	"gitee.com/jade-virtual/kunlun/component/log"
	"io"
	golog "log"
	"os"
	"time"

	"github.com/pkg/errors"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	gormlog "gorm.io/gorm/logger"
	"gorm.io/plugin/dbresolver"
)

type DB struct {
	*gorm.DB
}

type BaseMysqlPO struct {
	ID         uint64    `json:"id" gorm:"primarykey"`
	CreateTime time.Time `json:"create_time" gorm:"type:timestamptz;default:CURRENT_TIMESTAMP;not null"`
	UpdateTime time.Time `json:"update_time" gorm:"type:timestamptz;default:CURRENT_TIMESTAMP;not null"`
	Delete     bool      `json:"is_delete"` // 是否删除
}

func (m *DB) Health() bool {
	var probe struct{}
	if err := m.Raw("select 1").Scan(&probe).Error; err != nil {
		log.Error("database maybe dead ", log.Errval(err))
		return false
	}
	return true

}

func (m *DB) WithContext(ctx context.Context) *gorm.DB {
	d := m.DB.WithContext(ctx)

	return d
}

func New(c Config, opts ...Option) (*DB, error) {
	o := NewOptions(opts...)

	if c.Charset == "" {
		c.Charset = "utf8"
	}
	var (
		gormDb *gorm.DB
		err    error
	)

	switch c.Dialect {
	case MySQL:
		gormDb, err = newMysqlDB(c, o)
	case Postgres:
		gormDb, err = newPgDB(c, o)
	case SQLITE3:
		gormDb, err = gorm.Open(sqlite.Open(c.Address), &gorm.Config{})
	default:
		err = errors.New(string(c.Dialect) + " is not supported")
	}

	if err != nil {
		return nil, err
	}

	return &DB{DB: gormDb}, nil

}

var logWriter io.Writer

func SetLogOutput(w io.Writer) {
	logWriter = w
}

func newLog() gormlog.Interface {
	if logWriter == nil {
		logWriter = os.Stdout
	}
	return gormlog.New(golog.New(logWriter, "\r\n", golog.LstdFlags), gormlog.Config{
		SlowThreshold: 200 * time.Millisecond,
		LogLevel:      gormlog.Warn,
		Colorful:      true,
	})
}

func newDB(dbname, pDSN string, sDSN []string, opts Options, getDialector func(dsn string) gorm.Dialector) (*gorm.DB, error) {
	lvl := logger.Silent

	if opts.LogMode {
		lvl = logger.Info
	}

	gormDb, err := gorm.Open(getDialector(pDSN), &gorm.Config{
		Logger: newLog().LogMode(lvl),
	})

	if err != nil {
		return nil, err
	}

	if opts.LogMode {
		gormDb.Logger.LogMode(logger.Info)
	}

	if db, err := gormDb.DB(); err != nil {
		db.SetMaxIdleConns(opts.MaxIdleConn)
		db.SetMaxOpenConns(opts.MaxOpenConn)
	}

	if len(sDSN) == 0 {
		return gormDb, nil
	}

	rs := make([]gorm.Dialector, len(sDSN))

	for i, dsn := range sDSN {
		rs[i] = getDialector(dsn)
	}

	_ = gormDb.Use(dbresolver.Register(dbresolver.Config{
		Replicas: rs,
	}).SetMaxIdleConns(opts.MaxIdleConn).SetMaxOpenConns(opts.MaxOpenConn))

	return gormDb, nil
}

func newMysqlDB(c Config, opts Options) (*gorm.DB, error) {
	pDSN, sDSN := ParseMysqlDSN(c, opts)

	return newDB(c.DBName, pDSN, sDSN, opts, mysql.Open)
}

func newPgDB(c Config, opts Options) (*gorm.DB, error) {
	pDSN, sDSN := ParsePgDSN(c, opts)

	return newDB(c.DBName, pDSN, sDSN, opts, func(dsn string) gorm.Dialector {
		var conf = postgres.Config{
			DSN: dsn,
			// PreferSimpleProtocol 默认开启，否则数据结构变更将导致线上服务不可用。
			// 开启将导致 []byte 类型基础的，JSON 字段存储可能有问题，需要改为 string 基础的，详见：
			// https://github.com/go-gorm/datatypes/pull/11
			// https://git.5th.im/lb-public/gear/-/releases/v1.1.8
			PreferSimpleProtocol: true,
		}
		if opts.DisableSimpleProtocol {
			conf.PreferSimpleProtocol = false
		}

		return postgres.New(conf)
	})
}
